package com.atguigu.aiproject.service.impl;

import com.atguigu.aiproject.entity.Result;
import com.atguigu.aiproject.exception.BusinessException;
import com.atguigu.aiproject.mapper.ImgMapper;
import com.atguigu.aiproject.entity.Img;
import com.atguigu.aiproject.service.ImgService;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.List;

@Service
public class ImgServiceImpl extends ServiceImpl<ImgMapper, Img> implements ImgService {
    private final List<String> models;
    private final RestTemplate restTemplate;
    private final String url;
    private final String apiKey;

    public ImgServiceImpl(RestTemplate restTemplate,
                            @Value("${ai.image.model2.url}") String url,
                            @Value("${spring.ai.openai.api-key}") String apiKey) {
        this.restTemplate = restTemplate;
        this.url = url;
        this.apiKey = apiKey;
        this.models = List.of(
                "wanx2.0-t2i-turbo",
                "wanx-v1",
                "wanx2.1-t2i-turbo",
                "wanx2.1-t2i-plus",
                "wan2.2-t2i-flash",
                "wan2.2-t2i-plus",
                "wan2.5-t2i-preview"
        );
    }

    @Override
    public List<String> getModelsList() {
        return models;
    }

    @Override
    public String createImages(String prompt, int num) {
        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.set("X-DashScope-Async", "enable");
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", "Bearer " + apiKey);

        String requestBody = """
                {
                    "model": "%s",
                    "input": {
                        "prompt": "%s"
                    },
                    "parameters": {
                        "size": "1024*1024",
                        "n": %d
                    }
                }
                """.formatted(models.get(0), prompt, num);

        // 创建请求实体
        HttpEntity<String> requestEntity = new HttpEntity<>(requestBody, headers);

        // 发送POST请求并返回结果
        ResponseEntity<String> response = restTemplate.postForEntity(
                "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis",
                requestEntity,
                String.class
        );

        // 解析JSON提取task_id
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            JsonNode rootNode = objectMapper.readTree(response.getBody());

            // 获取task_id
            JsonNode taskIdNode = rootNode.path("output").path("task_id");
            String taskId = taskIdNode.asText();
            return taskId;
        } catch (Exception e) {
            return response.getBody();
        }
    }

    @Override
    public List<String> getImageStatus(String taskId) {
        // 创建请求头
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + apiKey);
        HttpEntity<String> requestEntity = new HttpEntity<>(headers);

        // 发送GET请求并返回结果
        String requestUrl = "https://dashscope.aliyuncs.com/api/v1/tasks/" + taskId;
        ResponseEntity<String> response = restTemplate.exchange(
                requestUrl,
                HttpMethod.GET,
                requestEntity,
                String.class
        );

        // 解析JSON提取图片URL数组
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            JsonNode rootNode = objectMapper.readTree(response.getBody());

            // 获取results数组
            JsonNode resultsNode = rootNode.path("output").path("results");

            // 创建图片URL列表
            List<String> imageUrls = new ArrayList<>();

            // 遍历results数组，提取每个url
            if (resultsNode.isArray()) {
                for (JsonNode result : resultsNode) {
                    JsonNode urlNode = result.path("url");
                    if (urlNode != null && !urlNode.asText().isEmpty()) {
                        imageUrls.add(urlNode.asText());
                    }
                }
            }
            return imageUrls;
        } catch (Exception e) {
            throw new BusinessException("获取图片失败");
        }
    }

    @Override
    public String createSingleImage(String prompt) {
        //1. 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", "Bearer " + apiKey);

        //2.设置请求体
        String requestBody = """
                   {
                              "model": "qwen-image",
                              "input": {
                                  "messages": [
                                      {
                                          "role": "user",
                                          "content": [
                                              {
                                                  "text": "%s"
                                              }
                                          ]
                                      }
                                  ]
                              },
                              "parameters": {
                                  "negative_prompt": "",
                                  "watermark": true,
                                  "size": "1328*1328",
                                  "prompt_extend":true
                              }
                   }
                """.formatted(prompt);

        // 创建请求实体
        HttpEntity<String> requestEntity = new HttpEntity<>(requestBody, headers);

        // 发送POST请求并返回结果
        ResponseEntity<String> response = restTemplate.postForEntity(
                url,
                requestEntity,
                String.class
        );

        // 解析JSON提取图片URL
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            JsonNode rootNode = objectMapper.readTree(response.getBody());
            JsonNode imageNode = rootNode.path("output")
                    .path("choices")
                    .get(0)
                    .path("message")
                    .path("content")
                    .get(0)
                    .path("image");

            String imageUrl = imageNode.asText();
            return imageUrl;
        } catch (Exception e) {
            // 如果解析失败，返回原始响应
            return response.getBody();
        }
    }
}
